package com.baishancloud.log

import org.apache.flink.api.common.state.{ListState, ListStateDescriptor}
import org.apache.flink.api.common.typeinfo.{TypeHint, TypeInformation}
import org.apache.flink.configuration.Configuration
import org.apache.flink.runtime.state.{FunctionInitializationContext, FunctionSnapshotContext}
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction
import org.apache.flink.streaming.api.functions.ProcessFunction
import org.apache.flink.util.Collector

import java.util.{Timer, TimerTask}
import scala.collection.mutable

/**
 * miniBatch实现，在keyBy之前先对数据进行一部分聚合，以减少下游数据量，进而减少数据倾斜造成的吞吐量量<br>
 * 默认为处理10w条数据，或者是处理10s之后输出一次中间聚合结果
 *
 * @author ziqiang.wang 
 * @date 2021/12/31 10:40
 */
class MiniBatchProcess extends ProcessFunction[MiguICTAuditLog, MiguICTAuditLog] with CheckpointedFunction {

  /**
   * 保留checkpoint快照时的中间结果状态
   */
  var listState: ListState[(MiguICTAuditLogKey, MiguICTAuditLogValue)] = _
  /**
   * 保存中间计算结果
   */
  val middleResult: mutable.Map[MiguICTAuditLogKey, MiguICTAuditLogValue] = mutable.Map[MiguICTAuditLogKey, MiguICTAuditLogValue]()
  /**
   * 输出阈值，处理多少条数据之后，输出所有的中间结果
   */
  val outThread: Long = 100000
  /**
   * 当前处理数据量
   */
  var count: Long = 0
  /**
   * 上次输出数据时间
   */
  var lastMilliSecond: Long = System.currentTimeMillis();
  /**
   * 当前处理时间，每秒更新一次
   */
  var currentMilliSecond: Long = lastMilliSecond
  /**
   * 输出时间间隔，固定10s
   */
  val outDuration: Long = 10 * 1000

  override def open(parameters: Configuration): Unit = {
    new Timer("定时更新当前处理时间", true).scheduleAtFixedRate(new TimerTask {
      override def run(): Unit = currentMilliSecond = System.currentTimeMillis()
    }, 1000, 1000)
  }

  override def processElement(value: MiguICTAuditLog, ctx: ProcessFunction[MiguICTAuditLog, MiguICTAuditLog]#Context, out: Collector[MiguICTAuditLog]): Unit = {
    val key: MiguICTAuditLogKey = value.getKey
    if (middleResult.contains(key)) {
      middleResult.put(key, MiguICTAuditLogValue(middleResult(key).line + value.line, middleResult(key).traffic + value.traffic))
    } else {
      middleResult.put(key, value.getValue)
    }
    count += 1
    if (count % outThread == 0 || currentMilliSecond - lastMilliSecond >= outDuration) {
      //  处理数据量达到输出阈值，或者是到达输出时间，则输出所有中间结果，然后清空中间结果表
      lastMilliSecond = currentMilliSecond
      middleResult.foreach(entry => out.collect(MiguICTAuditLog(entry._1, entry._2)))
      middleResult.clear()
    }

  }

  /**
   * 做快照时执行，将中间结果数据写入listState状态
   */
  override def snapshotState(context: FunctionSnapshotContext): Unit = {
    listState.clear()
    middleResult.foreach(entry => listState.add(entry._1, entry._2))
  }

  /**
   * 初始化和恢复快照时执行<br>
   * 初始化listState，将快照数据恢复到中间结果表
   */
  override def initializeState(context: FunctionInitializationContext): Unit = {
    val descriptor: ListStateDescriptor[(MiguICTAuditLogKey, MiguICTAuditLogValue)] = new ListStateDescriptor[(MiguICTAuditLogKey, MiguICTAuditLogValue)]("中间计算结果", TypeInformation.of(new TypeHint[(MiguICTAuditLogKey, MiguICTAuditLogValue)] {}))
    listState = context.getOperatorStateStore.getListState(descriptor)
    if (context.isRestored) {
      middleResult.clear()
      listState.get().forEach(entry => middleResult.put(entry._1, entry._2))
    }
  }
}
